/*
 * The Hygon TDM function test driver
 *
 * Copyright (C) 2018 Hygon Info Technologies Ltd.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 */

#include <linux/module.h>
#include <linux/slab.h>
#include <linux/spinlock.h>
#include <linux/spinlock_types.h>
#include <linux/err.h>
#include <asm/current.h>
#include <linux/kthread.h>
#include <linux/list.h>
#include <linux/delay.h>
#include <linux/kfifo.h>
#include <crypto/hash.h>
#include <linux/vmalloc.h>
#include <linux/version.h>
#if LINUX_VERSION_CODE < KERNEL_VERSION(6,4,0)
#include <linux/psp-sev.h>
#else
#include <linux/psp.h>
#endif
#include "tdm_hygon.h"

#define ERR_TDM_COND(cond, err, lable, fmt, ...) do {\
	if (cond) {					\
		ret = err;				\
		pr_err("Err: %s: "fmt, __func__, ##__VA_ARGS__);\
		goto lable;				\
	} } while (0)

static int test_scene = 0;
module_param(test_scene, int, 0644);
MODULE_PARM_DESC(test_scene, "test for different scene");

#define REGISTER_TASK_CNT 3
#define TEST_VMEM_SIZE		(RANGE_CNT_MAX*PAGE_SIZE)
char *vbuf[REGISTER_TASK_CNT] = {0};
char *kbuf[REGISTER_TASK_CNT][RANGE_CNT_MAX] = {0};
struct addr_range_info *addr_range_list[REGISTER_TASK_CNT] = {0};
struct measure_data *mdata_list[REGISTER_TASK_CNT] = {0};
struct authcode_2b *authcode_list[REGISTER_TASK_CNT] = {0};
int assign_task_id[REGISTER_TASK_CNT] = {0};

/*Calculates the hash value of the memory segment */
static int calc_task_kbuf_hash(int index, uint8_t *hash)
{
	int ret = 0;
	int i = 0;
	struct crypto_shash *shash = NULL;

	shash = crypto_alloc_shash("sm3", 0, 0);
	if (IS_ERR(shash)) {
		ret = PTR_ERR(shash);
		return ret;
	}

	{
		SHASH_DESC_ON_STACK(sdesc, shash);

		sdesc->tfm = shash;

		ret = crypto_shash_init(sdesc);
		if (ret) {
			pr_err("crypto_shash_init failed\n");
			ret = -1;
			goto out;
		}

		if (test_scene == 3) {
			ret = crypto_shash_update(sdesc, vbuf[index], TEST_VMEM_SIZE);
			if (ret) {
				pr_err("crypto_shash_update failed\n");
				ret = -1;
				goto out;
			}
		}
		else {
			for (i = 0; i < RANGE_CNT_MAX; i++) {
				ret = crypto_shash_update(sdesc, kbuf[index][i], (i + 1)*64);
				if (ret) {
					pr_err("crypto_shash_update failed\n");
					ret = -1;
					goto out;
				}
			}
		}

		ret = crypto_shash_final(sdesc, hash);
		if (ret) {
			pr_err("crypto_shash_final failed\n");
			ret = -1;
			goto out;
		}
	}

out:
	crypto_free_shash(shash);
	return ret;
}

static int regi_callback_handler(uint32_t task_id) {
	pr_info("ALARM!\n");
	pr_info("Task:%d, corruption detected!\n", task_id);
	pr_info("Please check if it's intended, or your machine may be on danger!\n");

	return 0;
}

//Allocate authorized memory to global pointer
static int task_authcode_memory_allocate(uint16_t count)
{
	int i = 0;
	int ret = 0;
	struct authcode_2b *authcode_d;

	ERR_TDM_COND(count > REGISTER_TASK_CNT, -DYN_BEYOND_MAX, end, "authcode count %d allocate failed!", count);

	for (i = 0; i < count; i++) {
		authcode_d = kzalloc(sizeof(struct authcode_2b) + AUTHCODE_MAX, GFP_KERNEL);
		ERR_TDM_COND(!authcode_d, -DYN_ERR_MEM, free_mem, "authcode_2b kzalloc memory failed\n");

		authcode_d->len = AUTHCODE_MAX;
		authcode_list[i] = authcode_d;
	}

end:
	return ret;

free_mem:
	for (i = 0; i < count; i++) {
		kfree(authcode_list[i]);
	}

	return ret;
}

//Allocate range and measure_data memory to global pointer
static int task_range_data_memory_allocate(struct addr_range_info **range, struct measure_data **data, int index)
{
	int i = 0, j = 0;
	int ret = 0;
	char *buf;
	struct addr_range_info *addr_range;
	struct measure_data *mdata_ptr;

	//allocate RANGE_CNT_MAX block memory,every block memory is padded the value base it's index
	for (i = 0; i < RANGE_CNT_MAX; i++) {
		buf = kzalloc((i + 1)*64, GFP_KERNEL);
		ERR_TDM_COND(!buf, -DYN_ERR_MEM, free_kbuf, "kbuf kzalloc memory failed\n");

		for(j = 0; j < (i + 1)*64; j++) {
			buf[j] = i + index;
		}

		kbuf[index][i] = buf;
	}

	addr_range = kzalloc(sizeof(struct addr_range_info) + \
			RANGE_CNT_MAX*sizeof(struct addr_info), GFP_KERNEL);
	ERR_TDM_COND(!addr_range, -DYN_ERR_MEM, free_kbuf, "addr_range kzalloc memory failed\n");

	addr_range->count = RANGE_CNT_MAX;
	for (i = 0; i < RANGE_CNT_MAX; i++) {
		addr_range->addr[i].addr_start = __psp_pa(kbuf[index][i]);
		addr_range->addr[i].length = (i + 1)*64;
	}

	mdata_ptr = kzalloc(sizeof(struct measure_data), GFP_KERNEL);
	ERR_TDM_COND(!mdata_ptr, -DYN_ERR_MEM, free_addr_range, "mdata_ptr kzalloc memory failed\n");

	//Configuration used by default
	mdata_ptr->hash_algo = HASH_ALGO_SM3;
	mdata_ptr->period_ms = 0;
	mdata_ptr->pcr = 10 + index % 7;
	ret = calc_task_kbuf_hash(index, mdata_ptr->expected_measurement);
	ERR_TDM_COND(ret, -DYN_ERR_MEM, free_addr_range, "calc_task_kbuf_hash failed\n");

	*range = addr_range;
	*data = mdata_ptr;

	return ret;

free_addr_range:
	kfree(addr_range);

free_kbuf:
	for (i = 0; i < RANGE_CNT_MAX; i++) {
		kfree(kbuf[index][i]);
	}

	return ret;
}

//Allocate virtual memory range and measure_data memory to global pointer
static int task_virtual_range_data_memory_allocate(struct addr_range_info **range, struct measure_data **data, int index)
{
	int ret = 0;
	struct addr_range_info *addr_range;
	struct measure_data *mdata_ptr;

	vbuf[index] = vmalloc(TEST_VMEM_SIZE);
	if (!vbuf[index]) {
		pr_err("vmalloc failed\n");
		return -1;
	}

	addr_range = kzalloc(sizeof(struct addr_range_info) + sizeof(struct addr_info), GFP_KERNEL);
	ERR_TDM_COND(!addr_range, -DYN_ERR_MEM, free_vbuf, "addr_range kzalloc memory failed\n");

	addr_range->count = 1;
	addr_range->addr[0].addr_start = (uint64_t)vbuf[index];
	addr_range->addr[0].length = TEST_VMEM_SIZE;

	mdata_ptr = kzalloc(sizeof(struct measure_data), GFP_KERNEL);
	ERR_TDM_COND(!mdata_ptr, -DYN_ERR_MEM, free_addr_range, "mdata_ptr kzalloc memory failed\n");

	//Configuration used by default
	mdata_ptr->hash_algo = HASH_ALGO_SM3;
	mdata_ptr->period_ms = 0;
	mdata_ptr->pcr = 10 + index % 7;
	ret = calc_task_kbuf_hash(index, mdata_ptr->expected_measurement);
	ERR_TDM_COND(ret, -DYN_ERR_MEM, free_addr_range, "calc_task_kbuf_hash failed\n");

	*range = addr_range;
	*data = mdata_ptr;

	return ret;

free_addr_range:
	kfree(addr_range);

free_vbuf:
	vfree(vbuf[index]);

	return ret;
}

//free task's related memory
static void task_memory_free(void)
{
	int i = 0, j = 0;

	for (i = 0; i < REGISTER_TASK_CNT; i++) {
		kfree(mdata_list[i]);
		vfree(vbuf[i]);
	}

	for (i = 0; i < REGISTER_TASK_CNT; i++) {
		for (j = 0; j < RANGE_CNT_MAX; j++) {
			kfree(kbuf[i][j]);
		}
	}

	for (i = 0; i < REGISTER_TASK_CNT; i++) {
		kfree(addr_range_list[i]);
	}

	for (i = 0; i < REGISTER_TASK_CNT; i++) {
		kfree(authcode_list[i]);
	}
}

static void release_task_list(void) {
	int i = 0;
	int ret = 0;
	int task_status = 0;
	struct measure_status task_measure_status;

	for (i = 0; i < REGISTER_TASK_CNT; i++) {
		ret = psp_query_measure_status(assign_task_id[i], &task_measure_status);
		ERR_TDM_COND(ret, ret, end, "After querying assign_task_id: %d status query failed\n", assign_task_id[i]);

		task_status = task_measure_status.status;
		if (task_status != DYN_STOP) {
			pr_info("Call psp_startstop_measure_task to stop measuring for task: %d", assign_task_id[i]);
			task_status = psp_startstop_measure_task(assign_task_id[i], authcode_list[i], false);
			ERR_TDM_COND(task_status < 0, task_status, end, "task_id %d stop failed\n", assign_task_id[i]);
			msleep(40);
		}

		pr_info("Call psp_destroy_measure_task to destroy measuring service for task: %d", assign_task_id[i]);
		ret = psp_destroy_measure_task(assign_task_id[i],authcode_list[i]);
		ERR_TDM_COND(ret, ret, end, "assign_task_id[%d] status destroy failed\n", assign_task_id[i]);
	}
	return;
end:
	pr_err("release_task failed!\n");
}

static int __init _module_init(void)
{
	int i = 0;
	int ret = 0;
	int task_status = 0;
	struct measure_update_data update_data;
	struct measure_status task_measure_status;
	struct authcode_2b authcode_err;

	update_data.update_flag = MEASURE_UPDATE_ALGO | MEASURE_UPDATE_EXPECTED_MEASUREMENT;
	update_data.algo = HASH_ALGO_SHA1;
	memset(update_data.expected_measurement, 0, 32);

#if !IS_ENABLED(CONFIG_CRYPTO_SM3)
	pr_err("Please check whether CONFIG_CRYPTO_SM3 is configured for this verify!!!\n");
	return -1;
#endif

	pr_info("\n–-----Victim module: has %d blocks of data measured by PSP--------test_scene:%d\n", REGISTER_TASK_CNT, test_scene);

	for (i = 0; i < REGISTER_TASK_CNT; i++) {
		if (test_scene == 3)
			ret = task_virtual_range_data_memory_allocate(&addr_range_list[i], &mdata_list[i], i);
		else
			ret = task_range_data_memory_allocate(&addr_range_list[i], &mdata_list[i], i);
		ERR_TDM_COND(ret != 0, ret, mem_free, "range data memory allocate failed!");
	}

	ret = task_authcode_memory_allocate(REGISTER_TASK_CNT);
	ERR_TDM_COND(ret != 0, ret, mem_free, "authcode memory allocate failed!");

	for (i = 0; i < REGISTER_TASK_CNT; i++) {
		pr_info("Call psp_create_measure_task to request measuring service.");
		if (test_scene == 3)
			assign_task_id[i] = psp_create_measure_task(addr_range_list[i], mdata_list[i], TASK_CREATE_VADDR | TASK_SUPPORT_VPCR, authcode_list[i]);
		else if (test_scene == 4)
			assign_task_id[i] = psp_create_measure_task(addr_range_list[i], mdata_list[i], TASK_SUPPORT_VPCR | TASK_EXCEPTION_CRASH, authcode_list[i]);
		else
			assign_task_id[i] = psp_create_measure_task(addr_range_list[i], mdata_list[i], TASK_SUPPORT_VPCR, authcode_list[i]);
		ERR_TDM_COND(assign_task_id[i] < 0, assign_task_id[i], mem_free, "task_id: 0x%x,create measurement task failed!", assign_task_id[i]);

		pr_info("Call psp_register_measure_exception_handler to register measuring exception function for task: %d", assign_task_id[i]);
		ret = psp_register_measure_exception_handler(assign_task_id[i], authcode_list[i], regi_callback_handler);
		ERR_TDM_COND(ret, ret, release_task, "task_id %d callback function register failed\n", assign_task_id[i]);

		pr_info("Call psp_startstop_measure_task to start measuring for task: %d", assign_task_id[i]);
		task_status = psp_startstop_measure_task(assign_task_id[i], authcode_list[i], true);
		ERR_TDM_COND(task_status < 0, task_status, release_task, "task_id %d start failed\n", assign_task_id[i]);

		msleep(40);
		ret = psp_query_measure_status(assign_task_id[i], &task_measure_status);
		ERR_TDM_COND(ret, ret, release_task, "task: %d status query failed\n", assign_task_id[i]);
	}

	switch (test_scene) {
	//scene0:
	//create->register->start->query
	case 0:
	//scene3:
	//vaddr create->register->start->query
	case 3:
		break;

	//scene1:
	//create->register->start->query->stop->update->start
	case 1:
	//scene4:
	//create->register->start->query->stop->update->start->exception hang
	case 4:
		msleep(100);
		for (i = 0; i < REGISTER_TASK_CNT; i++) {
			ret = psp_query_measure_status(assign_task_id[i], &task_measure_status);
			ERR_TDM_COND(ret, ret, release_task, "After querying assign_task_id: %d status query failed\n", assign_task_id[i]);

			task_status = task_measure_status.status;
			if (task_measure_status.status != DYN_STOP) {
				pr_info("Call psp_startstop_measure_task to stop measuring for task: %d", assign_task_id[i]);
				task_status = psp_startstop_measure_task(assign_task_id[i], authcode_list[i], false);
				ERR_TDM_COND(task_status < 0, task_status, release_task, "task_id %d stop failed\n", assign_task_id[i]);
				msleep(40);
			}

			pr_info("Call psp_update_measure_task to update measuring for task: %d", assign_task_id[i]);
			ret = psp_update_measure_task(assign_task_id[i], authcode_list[i], &update_data);
			ERR_TDM_COND(ret, ret, release_task, "assign_task_id: %d update failed\n", assign_task_id[i]);

			pr_info("Call psp_startstop_measure_task to start measuring for task: %d", assign_task_id[i]);
			task_status = psp_startstop_measure_task(assign_task_id[i], authcode_list[i], true);
			ERR_TDM_COND(task_status < 0, task_status, release_task, "task_id %d start failed\n", assign_task_id[i]);
		}

		break;

		//scene2:
		//create->register->start->error_auth->da->check da->check recovery->stop
	case 2:
		authcode_err.len = AUTHCODE_MAX;
		for (i = 0; i < REGISTER_TASK_CNT; i++) {
			pr_info("Call psp_startstop_measure_task in scene2 to stop measuring for task: %d", assign_task_id[i]);
			task_status = psp_startstop_measure_task(assign_task_id[i], &authcode_err, false);
			ERR_TDM_COND(task_status != -DYN_AUTH_FAIL, DYN_ERR_API, release_task, "stop task_id: %d  failed\n", assign_task_id[i]);
			msleep(1000);
			task_status = psp_startstop_measure_task(assign_task_id[i], &authcode_err, false);
			ERR_TDM_COND(task_status != -DYN_AUTH_FAIL, DYN_ERR_API, release_task, "stop task_id: %d  failed\n", assign_task_id[i]);
			task_status = psp_startstop_measure_task(assign_task_id[i], authcode_list[i], false);
			ERR_TDM_COND(task_status != -DYN_DA_PERIOD, DYN_ERR_API, release_task, "stop task_id: %d  failed\n", assign_task_id[i]);
			task_status = psp_startstop_measure_task(assign_task_id[i], authcode_list[i], false);
			ERR_TDM_COND(task_status != -DYN_DA_PERIOD, DYN_ERR_API, release_task, "stop task_id: %d  failed\n", assign_task_id[i]);
			msleep(2000);
			task_status = psp_startstop_measure_task(assign_task_id[i], authcode_list[i], false);
			ERR_TDM_COND(task_status < 0, DYN_ERR_API, release_task, "stop task_id: %d  failed\n", assign_task_id[i]);
		}
		break;

	default:
		pr_err("%s test_scene : %d is invalid \n", __func__, test_scene);
		ret = -1;
		break;
	}

	return ret;

release_task:
	release_task_list();
	pr_info("\n–----------------------end----------------------------\n");
	return ret;

mem_free:
	task_memory_free();
	return ret;
}

static void __exit _module_exit(void)
{
	release_task_list();
	task_memory_free();
	pr_info("\n–----------------------end----------------------------\n");
}

MODULE_AUTHOR("niuyongwen@hygon.cn");
MODULE_LICENSE("GPL");
MODULE_VERSION("0.7");
MODULE_DESCRIPTION("Functional testing demo for TDM");

module_init(_module_init);
module_exit(_module_exit);
